class ImageCluster
{
public:
    void Reset(const Bitmap &bmp, const Vec2i &seed);
    void AddCoord(const Bitmap &bmp, const Vec2i &coord);
    Vec2i MassCentroid(const Vec2i &dimensions);
    Vec2i ColorCentroid(const Bitmap &bmp, const Vec2i &dimensions);
    double AssignmentError(const Bitmap &bmp, const Vec2i &coord);

    __forceinline const Vec2i& seed()
    {
        return _seed;
    }
    __forceinline const Vector<Vec2i>& coords()
    {
        return _coords;
    }

private:
    Vec2i _seed;
    Vector<Vec2i> _coords;
};

class ImageSegmenter
{
public:
    struct QueueEntry
    {
        double priority;
        Vec2i coord;
        UINT clusterIndex;
    };

    void Segment(const Bitmap &bmp, UINT clusterCount, UINT iterationCount, Grid<UINT> &clusterIDs);
    static void DrawClusterIDs(const Grid<UINT> &clusterIDs, Bitmap &bmp);
    static void DrawClusterColors(const Bitmap &inputBmp, const Grid<UINT> &clusterIDs, Bitmap &outputBmp);

private:
    void InitializeClusters(const Bitmap &bmp);
    void AssignPixel(const Bitmap &bmp, const Vec2i &coord, UINT clusterIndex);
    void GrowClusters(const Bitmap &bmp);
    void RecenterClusters(const Bitmap &bmp);
    
    UINT _clusterSizeCutoff;
    Vec2i _dimensions;
    Vector<ImageCluster> _clusters;
    priority_queue<QueueEntry> _queue;
    Grid<UINT> _assignments;
};

__forceinline bool operator < (const ImageSegmenter::QueueEntry &a, const ImageSegmenter::QueueEntry &b)
{
    return (a.priority < b.priority);
}
